import numpy as np
import matplotlib.pyplot as plt

#定义函数y=x0*x0+x1*x1
def function_2(x):
    return x[0]**2 + x[1]**2

#定义计算梯度的函数
def numerical_gradient(f, x):
    h = 1e-4  #0.0001
    #生成和x形状相同的数组
    grad = np.zeros_like(x)

    for idx in range(x.size):
        tmp_val = x[idx]
        #计算f(x+h)
        x[idx] = tmp_val + h
        fxh1 = f(x)

        #计算法f(x-h)
        x[idx] = tmp_val - h
        fxh2 = f(x)

        grad[idx] = (fxh1 - fxh2) / (2*h)
        #还原值
        x[idx] = tmp_val
    return grad

#定义梯度下降法:参数f是要进行最优化的函数，init_x是初始值，lr是学习率,step_num是梯度发的重复次数
def gradient_descent(f, init_x, lr=0.01, step_num=100):
    x = init_x
    x_history = []

    for i in range(step_num):
        x_history.append( x.copy() )
        #求函数的梯度
        grad = numerical_gradient(f, x)
        x -= lr * grad
    x_history.append( x.copy() )
    return x, np.array(x_history)

init_x = np.array([-3.0, 4.0])
z, x_history = gradient_descent(function_2, init_x=init_x, lr=0.1, step_num=100)
print('学习率lr=0.1时： ' + str(z))
print(x_history)
print(init_x)

#梯度下降算法过程可视化
plt.plot([-5, 5], [0, 0], '--b')  #绘制x轴水平线
plt.plot([0, 0], [-5, 5], '--b')  #绘制y轴垂直线
plt.plot(x_history[:, 0], x_history[:, 1], 'o')
#plt.xlim(-3.5, 3.5)
#plt.ylim(-4.5, 4.5)
plt.xlabel('X0')
plt.ylabel('X1')
plt.show()

#比较不同学习率，学习率过大或者过小都无法得到好的结果
#学习率过大的例子：lr=10
init_x = np.array([-3.0, 4.0])
z, x_history = gradient_descent(function_2, init_x=init_x, lr=10, step_num=100)
print('学习率lr=10时： ' + str(z))
print(init_x)
#学习率过小的例子：lr=1e-10
init_x = np.array([-3.0, 4.0])
z, x_history = gradient_descent(function_2, init_x=init_x, lr=1e-10, step_num=100)
print('学习率lr=1e-10时： ' + str(z))
print(init_x)
